{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Multi-Agent System\n",
    "\n",
    "In this notebook, we showcase how the toolkit can be used to use a mixture of inbuilt tools and agents, as well as custom tools and workflows.\n",
    "\n",
    "We created a simple [mixture-of-agents](./retail_sales_agent/) that serves as an assistant in retail sales. \n",
    "> **Note**: *This is just an example agent system that uses dummy data. The intention is to demonstrate some of the capabilities of this toolkit and how a new user can get familiar with it.* \n",
    "\n",
    "This agent system has:\n",
    "1) A **supervisor** agent that routes incoming requests to the downstream agent expert\n",
    "2) A **data insight** agent that is a tool-calling agent capable of answering questions about sales data\n",
    "3) A **RAG agent** that is capable of answering questions about products using context from a product catalog\n",
    "4) A **data visualization** agent that is capable of plotting graphs and trends\n",
    "\n",
    "We demonstrate the following capabilities:\n",
    "- RAG\n",
    "- Multi-framework support\n",
    "- Human-in-the-Loop\n",
    "- Multi-agent support\n",
    "\n",
    "For more capabilities, refer to the `examples` directory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note**: \n",
    "> All source code for this example can be found at [./retail_sales_agent](./retail_sales_agent/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key\n",
    "\n",
    "if \"TAVILY_API_KEY\" not in os.environ:\n",
    "    tavily_api_key = getpass.getpass(\"Enter your Tavily API key: \")\n",
    "    os.environ[\"TAVILY_API_KEY\"] = tavily_api_key\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    openai_api_key = getpass.getpass(\"Enter your OpenAI API key: \")\n",
    "    os.environ[\"OPENAI_API_KEY\"] = openai_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a New Workflow for this Agent\n",
    "\n",
    "To recap, to create a new workflow for this mixture of agents, we need to use the `nat workflow create` sub-command which creates the necessary directory structure. \n",
    "\n",
    "> **Note**: You can create this directory structure manually as well.\n",
    "\n",
    "All new functions (tools and agents) that you want to be a part of this agent system can be created inside this directory for easier grouping of plugins. The only necessity for discovery by the toolkit is to import all new files/functions or simply define them in the `register.py` function.\n",
    "\n",
    "The example referenced in this notebook has already been created in the [retail_sales_agent](./retail_sales_agent/) uisng the following command:\n",
    "```bash\n",
    "nat workflow create --workflow-dir . retail_sales_agent\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start off simple, let's create a single agent that serves as a helpful assistant that can answer questions about the retail sales CSV data. It will call tools to fetch daily sales of a product, calculate total sales per day and detect any outliers in sales.\n",
    "\n",
    "**Function Creation**: All tools are created in [data_insight_tools.py](./retail_sales_agent/src/nat_retail_sales_agent/data_insight_tools.py). They each have a configuration object and the registered function.\n",
    "\n",
    "**Import the registered function**: Make sure to import the registered function in [register.py](./retail_sales_agent/src/nat_retail_sales_agent/register.py)\n",
    "\n",
    "**Create the YAML file**: For simplicity, we use the inbuilt react agent in the workflow and define the tools that should be made available to the agent. We also set the LLM to use. You can find the config file at [config.yml](./retail_sales_agent/configs/config.yml)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file retail_sales_agent/configs/config.yml --input \"How do laptop sales compare to phone sales?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some other test queries that can be run are:\n",
    "- \"What were the laptop sales on Feb 16th 2024?\"\n",
    "- \"What were the outliers in sales?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding a Retrieval Tool using Llamaindex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's add in a tool that is capable of performing retrieval of additional context to answer questions about products. It will use a vector store that stores details about products. We can create this agent using llama-index to demonstrate the framework-agnostic capability of the library. \n",
    "\n",
    "Refer to the code for the `product_catalog_rag` tool in [llama_index_rag_tool.py](./retail_sales_agent/src/nat_retail_sales_agent/llama_index_rag_tool.py). This can use a Milvus vector store for GPU-accelerated indexing. \n",
    "\n",
    "It requires the addition of an embedder section the [config_with_rag.yml](./retail_sales_agent/configs/config_with_rag.yml). This section follows a the same structure as the llms section and serves as a way to separate the embedding models from the LLM models. In our example, we are using the `nvidia/nv-embedqa-e5-v5` model.\n",
    "\n",
    "\n",
    "You can test this workflow with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file retail_sales_agent/configs/config_with_rag.yml --input \"What is the Ark S12 Ultra tablet and what are its specifications?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Agents and a Supervisor Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Building on the previous workflow, we can create an example that shows how to build a `react_agent` serving as a master orchestrator that routes queries to specialized `tool_calling_agent` or `react_agent` experts based on query content and agent descriptions. Further, it will exemplify how complete agent workflows can be wrapped and used as tools by other agents, enabling complex multi-agent orchestration.\n",
    "\n",
    "The full configuration file can be found at [config_multi_agent.yml](notebooks/retail_sales_agent/configs/config_multi_agent.yml)\n",
    "\n",
    "```yaml\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names: [data_analysis_agent, data_visualization_agent, rag_agent]\n",
    "  llm_name: supervisor_llm\n",
    "  verbose: true\n",
    "  handle_parsing_errors: true\n",
    "  max_retries: 2\n",
    "  system_prompt: |\n",
    "    Answer the following questions as best you can. You may communicate and collaborate with various experts to answer the questions:\n",
    "\n",
    "    {tools}\n",
    "\n",
    "    You may respond in one of two formats.\n",
    "    Use the following format exactly to communicate with an expert:\n",
    "\n",
    "    Question: the input question you must answer\n",
    "    Thought: you should always think about what to do\n",
    "    Action: the action to take, should be one of [{tool_names}]\n",
    "    Action Input: the input to the action (if there is no required input, include \"Action Input: None\")\n",
    "    Observation: wait for the expert to respond, do not assume the expert's response\n",
    "\n",
    "    ... (this Thought/Action/Action Input/Observation can repeat N times.)\n",
    "    Use the following format once you have the final answer:\n",
    "\n",
    "    Thought: I now know the final answer\n",
    "    Final Answer: the final answer to the original input question\n",
    "```\n",
    "\n",
    "The above workflow sections shows how a supervisor agent can be defined that behaves as the orchestrator and routes to downstream experts based on their function descriptions. The experts in this example are the previously created `data_analysis_agent` and two new agents - `rag_agent` created to handle RAG using the retrieval tool and `data_visualization_agent` to create plots and visualizations of data as requested by the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "nat run --config_file retail_sales_agent/configs/config_multi_agent.yml \\\n",
    "  --input \"What is the Ark S12 Ultra tablet and what are its specifications?\" \\\n",
    "  --input \"How do laptop sales compare to phone sales?\" \\\n",
    "  --input \"Plot average daily revenue\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom LangGraph Agent and Human-in-the-Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides using inbuilt agents in the workflows, we can also create custom agents using LangGraph or any other framework and bring them into a workflow. We demonstrate this by swapping out the `react_agent` used by the data visualization expert for a custom agent that has human-in-the-loop capability (utilizing a reusable plugin for HITL in the NeMo-Agent-Toolkit). The agent will ask the user whether they would like a summary of graph content.\n",
    "\n",
    "The code can be found in [data_visualization_agent.py](examples/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py)\n",
    "\n",
    "This agent has an agent node, a tools node, a node to accept human input and a summarizer node.\n",
    "\n",
    "Agent → generates tool calls → conditional_edge routes to tools\n",
    "\n",
    "Tools → execute → edge routes back to data_visualization_agent\n",
    "\n",
    "Agent → detects ToolMessage → creates summary AIMessage → conditional_edge routes to check_hitl_approval\n",
    "\n",
    "HITL → approval → conditional_edge routes to summarize or end\n",
    "\n",
    "\n",
    "#### Human-in-the-Loop Plugin\n",
    "\n",
    "This is enabled by leveraging a reusable plugin developed in the [examples/HITL/por_to_jiratickets](../HITL/por_to_jiratickets/) example. We can view the implementation in the [nat_por_to_jiratickets.hitl_approval_tool.py](../HITL/por_to_jiratickets/src/nat_por_to_jiratickets/hitl_approval_tool.py) file. The implementation is shown below:\n",
    "\n",
    "```python\n",
    "@register_function(config_type=HITLApprovalFnConfig)\n",
    "async def hitl_approval_function(config: HITLApprovalFnConfig, builder: Builder):\n",
    "\n",
    "    import re\n",
    "\n",
    "    prompt = f\"{config.prompt} Please confirm if you would like to proceed. Respond with 'yes' or 'no'.\"\n",
    "\n",
    "    async def _arun(unused: str = \"\") -> bool:\n",
    "\n",
    "        nat_context = Context.get()\n",
    "        user_input_manager = nat_context.user_interaction_manager\n",
    "\n",
    "        human_prompt_text = HumanPromptText(text=prompt, required=True, placeholder=\"<your response here>\")\n",
    "        response: InteractionResponse = await user_input_manager.prompt_user_input(human_prompt_text)\n",
    "        response_str = response.content.text.lower()  # type: ignore\n",
    "        selected_option = re.search(r'\\b(yes)\\b', response_str)\n",
    "\n",
    "        if selected_option:\n",
    "            return True\n",
    "        return False\n",
    "        # Rest of the function\n",
    "```\n",
    "\n",
    "As we see above, requesting user input using NeMo Agent toolkit is straightforward. We can use the user_input_manager to prompt the user for input. The user's response is then processed to determine the next steps in the workflow. This can occur in any tool or function in the workflow, allowing for dynamic interaction with the user as needed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the new workflow using the following command.\n",
    "\n",
    ">**Note**: This command needs to be run in a terminal since it requires accepting human input. Please open a terminal and run this command.\n",
    "\n",
    "```bash\n",
    "nat run --config_file retail_sales_agent/configs/config_multi_agent_hitl.yml --input \"Plot average daily revenue\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next Steps:\n",
    "\n",
    "The above feature examples are not exhaustive. The NeMo-Agent-Toolkit supports a continuously expanding list of features like [long-term memory support](../frameworks/semantic_kernel_demo) through partner integrations, [Model Context Protocol compatibility](../MCP/simple_calculator_mcp), a [demo chat UI](examples/UI), [custom API routes](../front_ends/simple_calculator_custom_routes) and so on. Please refer to the [examples](../) directory for a full catalog of examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To package and distribute your agent, the process is straightforward and follows standard Python `pyproject.toml` packaging steps. Refer to [this documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/extend/sharing-components.html) for a more detailed guide.\n",
    "\n",
    "Make sure to include all necessary NeMo Agent toolkit dependencies in the `pyproject.toml` as well as entrypoints.\n",
    "\n",
    "You can use the `nat info components` to discover the dependencies that need to be included in the `pyproject.toml`.\n",
    "\n",
    "Then you can either publish your package to a remote registry, build a wheel package for distribution, or share the source code."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
